# db.py (异步版本)
import aiomysql
from aiomysql.cursors import DictCursor
from config import Config

class DB:
    def __init__(self, config: Config):
        self.config = config
        self.pool = None  # 使用连接池

    async def connect(self):
        """创建连接池"""
        self.pool = await aiomysql.create_pool(
            host=self.config.DB_HOST,
            port=self.config.DB_PORT,
            user=self.config.DB_USERNAME,
            password=self.config.DB_PASSWORD,
            db=self.config.DB_DATABASE,
            cursorclass=DictCursor,
            autocommit=True,  # 启用自动提交
            minsize=1,  # 最小连接数
            maxsize=10   # 最大连接数
        )

    async def fetch_all(self, sql, params=None):
        if self.pool is None:
            await self.connect()
        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                await cursor.execute(sql, params)
                return await cursor.fetchall()

    async def fetch_one(self, sql, params=None):
        if self.pool is None:
            await self.connect()
        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                await cursor.execute(sql, params)
                return await cursor.fetchone()

    async def execute(self, sql, params=None):
        if self.pool is None:
            await self.connect()
        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                await cursor.execute(sql, params)
                return cursor.lastrowid

    async def insert_many(self, sql, values):
        if self.pool is None:
            await self.connect()
        async with self.pool.acquire() as conn:
            async with conn.cursor() as cursor:
                await cursor.executemany(sql, values)
                return cursor.rowcount

    async def close(self):
        """关闭连接池"""
        if self.pool:
            self.pool.close()
            await self.pool.wait_closed()